/*
 *    Copyright 2010 Vanessa Williams
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *        http://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */

/*
 * FASTASequence.java
 * Created on Feb 13, 2010
 */

package bioformatica.sequence;

import bioformatica.util.CloseableIterator;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.util.*;
import java.util.logging.Logger;

/**
 * Represents a sequence in a FASTA file.
 *
 * This class does not read the entire sequence into memory. Instead it makes use of
 * a read buffer and dynamic indexing for fast random access.
 *
 * Use the system property "READ_BUFFER_SIZE" to set the buffer size. Default size is
 * 128K (131,072). Trade memory for speed or vice versa.
 *
 * @author vwilliams
 * 
 */
public class FASTASequence extends Sequence {

    private static final Logger LOG = Logger.getLogger(FASTASequence.class.getName());

    private File fastaFile;
    private RandomAccessFile raf;
    private List<String> header = new ArrayList<String>();
    private long headerLength;
    private Alphabet alphabet;

    private static final String DEFAULT_READ_BUFFER_SIZE = "131072"; // 128K
    private static int READ_BUFFER_SIZE;
    private byte[] buffer;
    private String bufferOverflow = null;
    private int bufPosition = -1;
    private int bufSize = 0;

    CloseableIterator<Byte> iterator;
    private boolean beginIteration = true;

    {
        String bufSizeStr = System.getProperty("READ_BUFFER_SIZE", DEFAULT_READ_BUFFER_SIZE);
        READ_BUFFER_SIZE = Integer.parseInt(bufSizeStr);
    }

    /**
     * Constructor
     * @param fastaFile file containing sequence in FASTA format
     * @throws FileNotFoundException
     * @throws IOException if not a valid FASTA file
     */
    public FASTASequence(File fastaFile) throws FileNotFoundException, IOException {
        this.fastaFile = fastaFile;
        this.raf = new RandomAccessFile(fastaFile, "r");
        if (!isValid()) {
            throw new IOException("File does not appear to be in FASTA format; header is missing.");
        }
        readHeader();
        buffer = new byte[READ_BUFFER_SIZE];
        this.iterator  = new FASTASequenceIterator(this);
        // try to guess the alphabet used: DNA, RNA, Amino Acids
        guessAlphabet();
    }

    /**
     * {@inheritDoc}
     */
    public CloseableIterator<Byte> iterator() {
        if (iterator.isOpen()) throw new IllegalStateException("Iterator is already in use; must be closed before re-use.");
        iterator.open();
        return iterator;
    }

    /**
     * Not yet implemented.
     * {@inheritDoc}
     */
    public byte[] getSubsequence(int startPosition, int endPosition) {
        // TODO:
        throw new UnsupportedOperationException("Method not implemented yet");
    }

    /**
     * @return  File header
     */
    public List<String> getHeader() {
        return this.header;
    }

    /**
     * {@inheritDoc}
     */
    public Alphabet getAlphabet() {
        return this.alphabet;
    }

    protected synchronized boolean hasNext() {

        if (bufPosition < 0) {
            if (!fillBuffer()) {
                beginIteration = true; // eof, so set flag to seek to start next time we fill buffer
                return false;
            }
        }
        return true;
    }

    protected synchronized Byte next() {
        Byte result = buffer[bufPosition++];
        if (bufPosition >= bufSize) {
            bufPosition = -1; // time to refill the buffer
            bufSize = 0;
        }
        return result;
    }

    private boolean isValid() throws IOException {
        raf.seek(0);
        boolean valid = true;
        byte[] firstChar = new byte[1];
        firstChar[0] = raf.readByte();
        if (!(new String(firstChar).equals(">"))) {
            valid = false;
        }
        raf.seek(0);
        return valid;
    }

    private void readHeader() throws IOException {
        long filePointer = 0;
        while (true) {
            String potentialHeaderLine = raf.readLine();
            if (potentialHeaderLine == null) {
                break;
            }
            if (potentialHeaderLine.startsWith(">")) {
                this.header.add(potentialHeaderLine);
                filePointer = raf.getFilePointer();
            }
            else {
                break;
            }
        }
        this.headerLength = filePointer;
    }

    private boolean fillBuffer() {

        int readLength = 0;

        // add any overflow that might have been left from the last read.
        if (bufferOverflow != null) {
            bufferAppend(bufferOverflow.getBytes(), bufferOverflow.length());
            readLength += bufferOverflow.length();
            bufferOverflow = null;
        }

        if (beginIteration) {
            // this is the beginning of an iterator's usage, so seek to start of sequence
            try {
                raf.seek(headerLength);
                beginIteration = false;
            } catch (IOException e) {
                LOG.warning("Unable to seek to start of sequence: " + e.getMessage());
            }
        }

        String read;
        boolean full = false;
        boolean eof = false;

        while (!full && !eof) {
            try {
                read = raf.readLine();
                if (read == null) {
                    eof = true;
                }
                else {
                    int readStrLength = read.length();
                    int newBufLength = readLength + readStrLength;
                    if (newBufLength >= READ_BUFFER_SIZE) {
                        full = true;
                        if (newBufLength > READ_BUFFER_SIZE) {
                            // shorten read so it fits in buffer
                            int room = READ_BUFFER_SIZE - readLength;
                            String savedFullRead = new String(read);
                            read = read.substring(0, room);
                            readStrLength = read.length();
                            // save the extra read substring for the next fill
                            bufferOverflow = savedFullRead.substring(room);
                        }
                    }
                    bufferAppend(read.getBytes(), readStrLength);
                    readLength += readStrLength;
                }
            } catch (IOException e) {
                LOG.warning("Exception reading FASTA file: " + e.getMessage());
            }
        }
        
        if (readLength != 0) {
            bufPosition = 0; // it's ok to start reading now
            return true;
        }
        else {
            return false;
        }
    }

    private void bufferAppend(byte[] toAppend, int length) {
        for (int i=0; i<length; i++) {
            buffer[bufSize++] = toAppend[i];
        }
    }

    private void guessAlphabet() throws IOException {
        
        long fileLength = raf.length();
        long[] seekPositions = new long[4];
        seekPositions[0] = headerLength;
        seekPositions[1] = fileLength/4;
        seekPositions[2] = fileLength/2;
        seekPositions[3] = fileLength/2 + fileLength/4;

        // read some sample data from 4 different positions in the file
        StringBuilder sb = new StringBuilder();
        for (int i=0; i<4; i++) {
            raf.seek(seekPositions[i]);
            for (int j=0; j<10; j++) {
                try {
                    sb.append(raf.readLine());
                } catch (IOException e) {
                    // ignore: we really don't care if we get fewer lines than we asked for
                }
            }
        }

        // scan the sample and determine the set size of its contents
        byte[] sampleBytes = sb.toString().getBytes();
        Set<Byte> charSet = new HashSet<Byte>();
        for (int i=0; i<sb.length(); i++) {
            charSet.add(sampleBytes[i]);
        }
        // use a heuristic to guess the alphabet based on the size and contents of the set
        int numElements = charSet.size();
        if (numElements <= 10) {
            // it's either RNA or DNA
            if (charSet.contains("U".getBytes()[0])) {
                this.alphabet = Alphabet.RNA;
            }
            else {
                this.alphabet = Alphabet.DNA;
            }
        }
        else {
            this.alphabet = Alphabet.PROTEIN;
        }
    }
}
